#!/usr/bin/env python3
"""
MyLangExtract 最终测试脚本
使用更简单的提示词和更直接的方法
"""
import os
import sys
import json
import time
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, asdict

try:
    import requests
except ImportError:
    print("❌ 缺少 requests 库")
    print("请运行: pip install requests")
    sys.exit(1)

@dataclass
class Extraction:
    """提取结果数据类"""
    extraction_class: str
    extraction_text: str
    attributes: Optional[Dict[str, str]] = None
    
    def __post_init__(self):
        if self.attributes is None:
            self.attributes = {}

def get_api_key() -> Optional[str]:
    """获取智谱清言 API 密钥"""
    return os.environ.get('ZHIPU_API_KEY')

def simple_extraction_test():
    """简单的提取测试"""
    print("🧪 MyLangExtract 最终功能测试")
    print("=" * 50)
    
    api_key = get_api_key()
    if not api_key:
        print("❌ 未配置 ZHIPU_API_KEY")
        print("请运行: export ZHIPU_API_KEY='your_api_key_here'")
        return False
    
    print("✅ 智谱清言 API 密钥已配置")
    
    # 构建简单的请求
    url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }
    
    # 使用更简单的 function calling
    schema = {
        "name": "extract_info",
        "description": "提取人物和职业信息",
        "parameters": {
            "type": "object",
            "properties": {
                "person": {
                    "type": "string",
                    "description": "人物姓名"
                },
                "job": {
                    "type": "string", 
                    "description": "职业"
                },
                "age": {
                    "type": "string",
                    "description": "年龄"
                }
            },
            "required": ["person", "job"]
        }
    }
    
    # 简化的提示词
    prompt = """请从以下文本中提取人物姓名、职业和年龄信息：

文本：李四是一名产品经理，今年30岁

请使用 extract_info 函数返回结果。"""
    
    data = {
        'model': 'glm-4',
        'messages': [
            {'role': 'user', 'content': prompt}
        ],
        'tools': [{
            'type': 'function',
            'function': schema
        }],
        'tool_choice': {'type': 'function', 'function': {'name': 'extract_info'}},
        'temperature': 0.1
    }
    
    try:
        print("🔄 正在调用智谱清言 API...")
        start_time = time.time()
        response = requests.post(url, headers=headers, json=data, timeout=30)
        end_time = time.time()
        
        print(f"📡 API 响应状态: {response.status_code}")
        
        if response.status_code == 200:
            result = response.json()
            print(f"📋 API 响应: {json.dumps(result, ensure_ascii=False, indent=2)}")
            
            # 解析结果
            if 'choices' in result and len(result['choices']) > 0:
                choice = result['choices'][0]
                
                if 'message' in choice and 'tool_calls' in choice['message']:
                    tool_calls = choice['message']['tool_calls']
                    if tool_calls and len(tool_calls) > 0:
                        tool_call = tool_calls[0]
                        if tool_call['function']['name'] == 'extract_info':
                            try:
                                arguments = json.loads(tool_call['function']['arguments'])
                                
                                print(f"✅ 提取成功！")
                                print(f"📊 处理时间: {end_time - start_time:.2f}秒")
                                print(f"🎯 提取结果:")
                                
                                for key, value in arguments.items():
                                    if value:
                                        print(f"  • {key}: {value}")
                                
                                # 创建标准化的提取结果
                                extractions = []
                                if arguments.get('person'):
                                    extractions.append(Extraction("person", arguments['person']))
                                if arguments.get('job'):
                                    extractions.append(Extraction("job", arguments['job']))
                                if arguments.get('age'):
                                    extractions.append(Extraction("age", arguments['age']))
                                
                                print(f"\n📝 标准化结果:")
                                for i, ext in enumerate(extractions, 1):
                                    print(f"  {i}. {ext.extraction_class}: '{ext.extraction_text}'")
                                
                                return True
                                
                            except json.JSONDecodeError as e:
                                print(f"❌ 解析结果失败: {e}")
                                return False
                
                print("❌ 未找到有效的 tool call 结果")
                return False
            
            print("❌ API 响应格式不正确")
            return False
        else:
            print(f"❌ API 请求失败: {response.status_code}")
            print(f"错误信息: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ API 调用异常: {e}")
        return False

def test_multiple_examples():
    """测试多个示例"""
    print("\n🔄 测试多个示例...")
    
    api_key = get_api_key()
    if not api_key:
        return False
    
    test_cases = [
        "王五是一名医生，今年35岁",
        "赵六在北京做设计师，28岁",
        "钱七是老师，工作了5年"
    ]
    
    url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }
    
    schema = {
        "name": "extract_info",
        "description": "提取人物和职业信息",
        "parameters": {
            "type": "object",
            "properties": {
                "person": {"type": "string", "description": "人物姓名"},
                "job": {"type": "string", "description": "职业"},
                "age": {"type": "string", "description": "年龄"},
                "location": {"type": "string", "description": "地点"}
            },
            "required": ["person"]
        }
    }
    
    success_count = 0
    
    for i, text in enumerate(test_cases, 1):
        print(f"\n测试 {i}: {text}")
        
        prompt = f"请从以下文本中提取人物信息：\n\n文本：{text}\n\n请使用 extract_info 函数返回结果。"
        
        data = {
            'model': 'glm-4',
            'messages': [{'role': 'user', 'content': prompt}],
            'tools': [{'type': 'function', 'function': schema}],
            'tool_choice': {'type': 'function', 'function': {'name': 'extract_info'}},
            'temperature': 0.1
        }
        
        try:
            response = requests.post(url, headers=headers, json=data, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                
                if ('choices' in result and len(result['choices']) > 0 and
                    'message' in result['choices'][0] and 
                    'tool_calls' in result['choices'][0]['message']):
                    
                    tool_calls = result['choices'][0]['message']['tool_calls']
                    if tool_calls and len(tool_calls) > 0:
                        arguments = json.loads(tool_calls[0]['function']['arguments'])
                        
                        print(f"  ✅ 成功提取:")
                        for key, value in arguments.items():
                            if value:
                                print(f"    {key}: {value}")
                        success_count += 1
                    else:
                        print(f"  ❌ 未找到提取结果")
                else:
                    print(f"  ❌ 响应格式错误")
            else:
                print(f"  ❌ API 错误: {response.status_code}")
                
        except Exception as e:
            print(f"  ❌ 异常: {e}")
    
    print(f"\n📊 批量测试结果: {success_count}/{len(test_cases)} 成功")
    return success_count > 0

def main():
    """主函数"""
    print("MyLangExtract 最终功能验证")
    print("=" * 50)
    
    # 基础功能测试
    if simple_extraction_test():
        print("\n🎉 基础功能测试通过！")
        
        # 多示例测试
        if test_multiple_examples():
            print("\n" + "=" * 50)
            print("🎉 所有测试通过！MyLangExtract 核心功能正常工作")
            print("\n✨ 您的工具已经可以:")
            print("  • 连接智谱清言 API")
            print("  • 执行结构化信息提取")
            print("  • 处理多种文本格式")
            print("  • 返回标准化的提取结果")
            
            print("\n📚 下一步:")
            print("  • 查看 QUICKSTART.md 了解完整用法")
            print("  • 修改此脚本适应您的具体需求")
            print("  • 集成到您的其他项目中")
            
            print("\n🔧 自定义使用:")
            print("  • 修改 schema 定义不同的提取字段")
            print("  • 调整提示词以适应不同的文本类型")
            print("  • 添加更多的提供商支持")
        else:
            print("\n⚠️  批量测试部分失败，但基础功能正常")
    else:
        print("\n❌ 基础功能测试失败")
        print("请检查:")
        print("  • API 密钥是否正确")
        print("  • 网络连接是否正常")
        print("  • 智谱清言服务是否可用")

if __name__ == "__main__":
    main()
